"""
融合跟踪器实现模块

该模块实现了MeanShift算法与Kalman滤波器相结合的目标跟踪器。
主要功能：
1. 融合两种算法的优势
2. 处理目标遮挡和快速运动
3. 自适应参数调整
4. 跟踪状态管理

作者：目标跟踪系统
日期：2024
"""

import cv2
import numpy as np
from typing import Tuple, Optional, List
from .meanshift import MeanShiftTracker
from .kalman import KalmanFilter


class FusionTracker:
    """
    融合跟踪器类
    
    结合MeanShift算法和Kalman滤波器的优势，实现鲁棒的目标跟踪。
    通过预测-搜索-验证-更新的流程，提高跟踪的准确性和稳定性。
    
    参数：
        similarity_threshold: 相似度阈值，用于判断跟踪是否成功
        max_lost_frames: 最大丢失帧数，超过则认为跟踪失败
        kalman_weight: Kalman滤波器权重，用于融合两种算法结果
    """
    
    def __init__(self, similarity_threshold: float = 0.3, max_lost_frames: int = 10, kalman_weight: float = 0.3):
        # 初始化子跟踪器
        self.meanshift_tracker = MeanShiftTracker()
        self.kalman_filter = KalmanFilter()
        
        # 跟踪参数
        self.similarity_threshold = similarity_threshold
        self.max_lost_frames = max_lost_frames
        self.kalman_weight = kalman_weight
        
        # 跟踪状态
        self.is_initialized = False
        self.current_bbox = None
        self.lost_frames = 0
        self.tracking_history = []
        
        # 性能统计
        self.frame_count = 0
        self.successful_tracks = 0
        
    def initialize(self, image: np.ndarray, bbox: Tuple[int, int, int, int]) -> None:
        """
        初始化跟踪器
        
        设置目标模型和初始状态，准备开始跟踪。
        
        参数：
            image: 初始帧图像，BGR格式
            bbox: 目标边界框 (x, y, width, height)
        """
        # 初始化MeanShift跟踪器
        self.meanshift_tracker.initialize_target(image, bbox)
        
        # 初始化Kalman滤波器
        center_x = bbox[0] + bbox[2] // 2
        center_y = bbox[1] + bbox[3] // 2
        self.kalman_filter.initialize((center_x, center_y))
        
        # 设置初始状态
        self.current_bbox = bbox
        self.is_initialized = True
        self.lost_frames = 0
        self.tracking_history = [(center_x, center_y)]
        
        # 重置统计信息
        self.frame_count = 0
        self.successful_tracks = 0
        
    def track(self, image: np.ndarray) -> Tuple[bool, Optional[Tuple[int, int, int, int]]]:
        """
        在新帧中跟踪目标
        
        使用融合算法跟踪目标，返回跟踪结果和状态。
        
        参数：
            image: 当前帧图像，BGR格式
            
        返回：
            (跟踪是否成功, 目标边界框)
        """
        if not self.is_initialized:
            raise ValueError("跟踪器未初始化，请先调用initialize方法")
            
        self.frame_count += 1
        
        # 步骤1：Kalman滤波器预测
        predicted_position = self.kalman_filter.predict()
        
        # 步骤2：构建搜索窗口
        search_window = self._get_search_window(predicted_position, image.shape)
        
        # 步骤3：MeanShift搜索
        meanshift_bbox = self.meanshift_tracker.track(image, search_window)
        
        # 步骤4：验证跟踪结果
        similarity = self.meanshift_tracker.get_similarity(image, meanshift_bbox)
        
        if similarity > self.similarity_threshold:
            # 跟踪成功
            self.lost_frames = 0
            self.successful_tracks += 1
            
            # 步骤5：融合结果
            fused_bbox = self._fuse_results(predicted_position, meanshift_bbox)
            
            # 步骤6：更新Kalman滤波器
            center_x = fused_bbox[0] + fused_bbox[2] // 2
            center_y = fused_bbox[1] + fused_bbox[3] // 2
            self.kalman_filter.update((center_x, center_y))
            
            # 更新状态
            self.current_bbox = fused_bbox
            self.tracking_history.append((center_x, center_y))
            
            # 自适应参数调整
            self._adaptive_parameter_adjustment(similarity)
            
            return True, fused_bbox
            
        else:
            # 跟踪失败
            self.lost_frames += 1
            
            if self.lost_frames < self.max_lost_frames:
                # 使用Kalman预测结果
                predicted_bbox = self._position_to_bbox(predicted_position)
                self.current_bbox = predicted_bbox
                return False, predicted_bbox
            else:
                # 跟踪完全失败
                return False, None
                
    def _get_search_window(self, predicted_position: Tuple[float, float], image_shape: Tuple[int, int, int]) -> Tuple[int, int, int, int]:
        """
        根据预测位置构建搜索窗口
        
        参数：
            predicted_position: 预测的目标中心位置
            image_shape: 图像尺寸 (height, width, channels)
            
        返回：
            搜索窗口 (x, y, width, height)
        """
        if self.current_bbox is None:
            return (0, 0, image_shape[1], image_shape[0])
            
        # 获取当前目标尺寸
        _, _, w, h = self.current_bbox
        
        # 根据不确定性调整搜索窗口大小
        uncertainty = self.kalman_filter.get_position_uncertainty()
        scale_factor = min(3.0, 1.5 + uncertainty / 50.0)
        
        search_w = int(w * scale_factor)
        search_h = int(h * scale_factor)
        
        # 计算搜索窗口位置
        search_x = max(0, int(predicted_position[0] - search_w // 2))
        search_y = max(0, int(predicted_position[1] - search_h // 2))
        
        # 确保搜索窗口在图像范围内
        search_w = min(search_w, image_shape[1] - search_x)
        search_h = min(search_h, image_shape[0] - search_y)
        
        return (search_x, search_y, search_w, search_h)
        
    def _fuse_results(self, predicted_position: Tuple[float, float], meanshift_bbox: Tuple[int, int, int, int]) -> Tuple[int, int, int, int]:
        """
        融合Kalman预测和MeanShift结果
        
        参数：
            predicted_position: Kalman预测位置
            meanshift_bbox: MeanShift跟踪结果
            
        返回：
            融合后的边界框
        """
        # 获取MeanShift结果的中心位置
        ms_center_x = meanshift_bbox[0] + meanshift_bbox[2] // 2
        ms_center_y = meanshift_bbox[1] + meanshift_bbox[3] // 2
        
        # 加权融合位置
        fused_x = (1 - self.kalman_weight) * ms_center_x + self.kalman_weight * predicted_position[0]
        fused_y = (1 - self.kalman_weight) * ms_center_y + self.kalman_weight * predicted_position[1]
        
        # 保持目标尺寸
        w, h = meanshift_bbox[2], meanshift_bbox[3]
        
        # 转换为边界框格式
        bbox_x = int(fused_x - w // 2)
        bbox_y = int(fused_y - h // 2)
        
        return (bbox_x, bbox_y, w, h)
        
    def _position_to_bbox(self, position: Tuple[float, float]) -> Tuple[int, int, int, int]:
        """
        将位置转换为边界框
        
        参数：
            position: 中心位置 (x, y)
            
        返回：
            边界框 (x, y, width, height)
        """
        if self.current_bbox is None:
            return (int(position[0]), int(position[1]), 50, 50)
            
        w, h = self.current_bbox[2], self.current_bbox[3]
        x = int(position[0] - w // 2)
        y = int(position[1] - h // 2)
        
        return (x, y, w, h)
        
    def _adaptive_parameter_adjustment(self, similarity: float) -> None:
        """
        自适应参数调整
        
        根据跟踪性能动态调整算法参数。
        
        参数：
            similarity: 当前帧的相似度
        """
        # 根据相似度调整Kalman权重
        if similarity > 0.7:
            # 高相似度，更信任MeanShift结果
            self.kalman_weight = max(0.1, self.kalman_weight - 0.01)
        elif similarity < 0.4:
            # 低相似度，更信任Kalman预测
            self.kalman_weight = min(0.5, self.kalman_weight + 0.01)
            
        # 根据跟踪历史调整过程噪声
        if len(self.tracking_history) > 5:
            recent_positions = self.tracking_history[-5:]
            velocity_variance = self._calculate_velocity_variance(recent_positions)
            
            # 动态调整过程噪声
            if velocity_variance > 100:
                self.kalman_filter.set_process_noise(0.1)
            else:
                self.kalman_filter.set_process_noise(0.01)
                
    def _calculate_velocity_variance(self, positions: List[Tuple[float, float]]) -> float:
        """
        计算速度方差
        
        参数：
            positions: 位置历史列表
            
        返回：
            速度方差
        """
        if len(positions) < 2:
            return 0.0
            
        velocities = []
        for i in range(1, len(positions)):
            vx = positions[i][0] - positions[i-1][0]
            vy = positions[i][1] - positions[i-1][1]
            velocities.append((vx, vy))
            
        if not velocities:
            return 0.0
            
        # 计算速度方差
        vx_mean = np.mean([v[0] for v in velocities])
        vy_mean = np.mean([v[1] for v in velocities])
        
        vx_var = np.var([v[0] for v in velocities])
        vy_var = np.var([v[1] for v in velocities])
        
        return vx_var + vy_var
        
    def get_tracking_info(self) -> dict:
        """
        获取跟踪信息
        
        返回：
            包含跟踪统计信息的字典
        """
        success_rate = self.successful_tracks / max(1, self.frame_count)
        
        return {
            'frame_count': self.frame_count,
            'successful_tracks': self.successful_tracks,
            'success_rate': success_rate,
            'lost_frames': self.lost_frames,
            'is_tracking': self.lost_frames < self.max_lost_frames,
            'current_bbox': self.current_bbox,
            'kalman_weight': self.kalman_weight,
            'similarity_threshold': self.similarity_threshold
        }
        
    def reset(self) -> None:
        """
        重置跟踪器状态
        
        清除所有跟踪信息，准备重新初始化。
        """
        self.is_initialized = False
        self.current_bbox = None
        self.lost_frames = 0
        self.tracking_history = []
        self.frame_count = 0
        self.successful_tracks = 0
        
        # 重置子跟踪器
        self.kalman_filter.reset()
        
    def set_parameters(self, similarity_threshold: Optional[float] = None, 
                      max_lost_frames: Optional[int] = None,
                      kalman_weight: Optional[float] = None) -> None:
        """
        设置跟踪参数
        
        参数：
            similarity_threshold: 相似度阈值
            max_lost_frames: 最大丢失帧数
            kalman_weight: Kalman权重
        """
        if similarity_threshold is not None:
            self.similarity_threshold = similarity_threshold
        if max_lost_frames is not None:
            self.max_lost_frames = max_lost_frames
        if kalman_weight is not None:
            self.kalman_weight = kalman_weight 